package com.chukun.flink.stream.state.broadcast;

import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.util.Collector;

import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * @author chukun
 * @version 1.0.0
 * @description 广播基本操作
 * @createTime 2022年05月16日 21:26:00
 */
public class BroadcastStateTemplate {

    /**
     * 定义原始数据的数据源
     */
    private static class CustomSource implements SourceFunction<Date> {

        @Override
        public void run(SourceContext<Date> context) throws Exception {
            while (true) {
                TimeUnit.SECONDS.sleep(1);
                context.collect(new Date());
            }
        }

        @Override
        public void cancel() {}
    }

    /**
     * 定义规则数据源
     */
    private static class RuleSource implements SourceFunction<Tuple2<Integer, String>> {

        // 定义一组日期格式化规则
        private String[] formats = new String[]{"yyyy-MM-dd HH:mm", "yyyy-MM-dd HH","yyyy-MM-dd","yyyy-MM", "yyyy"};

        @Override
        public void run(SourceContext<Tuple2<Integer, String>> context) throws Exception {
            while (true) {
                // 遍历日期格式化数组，每隔5s输出一组日期规则
                for (String format : formats) {
                    context.collect(new Tuple2<>(1, format));
                    TimeUnit.SECONDS.sleep(5);
                }
            }
        }

        @Override
        public void cancel() {}
    }

    /**
     * 获取广播中的数据对非广播中的数据进行校验
     */
    private static class CustomBroadcastProcessFunction extends BroadcastProcessFunction<Date, Tuple2<Integer,String>, Tuple2<String, String>> {


        private transient MapStateDescriptor<Integer,String> descriptor;

        @Override
        public void open(Configuration parameters) throws Exception {
            // 声明广播状态的名称及状态保存的值类型
            // 要与ruleStream.broadcast(stateDesc)中定义的状态描述符保持一致
            /**
             * {@link BroadcastStateTemplate#main(String[])}
             */
            descriptor = new MapStateDescriptor<>("broadcast-state", Integer.class, String.class);
        }

        // 每个非广播流中元素的处理方法
        @Override
        public void processElement(Date date, BroadcastProcessFunction<Date, Tuple2<Integer, String>, Tuple2<String, String>>.ReadOnlyContext readOnlyContext, Collector<Tuple2<String, String>> collector) throws Exception {
            String formatRule ="";
            for (Map.Entry<Integer, String> entry : readOnlyContext.getBroadcastState(descriptor).immutableEntries()) {
                // 判断广播状态中的key是否为1，为1代表是格式化的日期格式
                if (entry.getKey() == 1) {
                    formatRule = entry.getValue();
                }
            }
            if (StringUtils.isNotEmpty(formatRule)) {
                String originalDate = new SimpleDateFormat("yyyy-MM-dd HH;mm:ss").format(date);
                // 根据广播状态中的格式化日志规则对非广播中的元素进行日期格式化
                String formatDate = new SimpleDateFormat(formatRule).format(date);
                collector.collect(Tuple2.of("主数据流元素:" + originalDate, "应用规则后的格式：" + formatDate));
            }
        }

        // 每个广播流中元素的处理方法
        @Override
        public void processBroadcastElement(Tuple2<Integer, String> input, BroadcastProcessFunction<Date, Tuple2<Integer, String>, Tuple2<String, String>>.Context context, Collector<Tuple2<String, String>> collector) throws Exception {
            // 根据指定的状态描述符获取广播状态
            BroadcastState<Integer, String> broadcastState = context.getBroadcastState(descriptor);
            // 将新的规则存入广播状态中
            broadcastState.put(input.f0, input.f1);
            // 可选，输出规则
            collector.collect(new Tuple2<>("广播状态中新增元素", input.toString()));
        }
    }

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        // 设置全局并发度
        env.setParallelism(2);

        // 加载规则流数据
        DataStream<Tuple2<Integer, String>> ruleStream = env.addSource(new RuleSource());

        // 加载原始数据流
        DataStream<Date> dataStream = env.addSource(new CustomSource());

        // 声明广播状态的名称及状态保存的值类型，状态的key存储规则id，value存储规则的描述信息
        final MapStateDescriptor<Integer,String> stateDesc = new MapStateDescriptor<>("broadcast-state", Integer.class, String.class);

        // 将ruleStream数据流转换为一个具有广播状态的数据流
        BroadcastStream<Tuple2<Integer, String>> broadcastStream = ruleStream.broadcast(stateDesc);

        // 数据流调用connect方法连接广播流，生成BroadcastConnectedStream数据流中调用process方法来应用自定义的BroadcastProcessFunction对数据元素校验，转换
        DataStream<Tuple2<String, String>> result = dataStream.connect(broadcastStream).process(new CustomBroadcastProcessFunction());

        result.print("输出结果");

        env.execute("BroadcastStateTemplate");

    }
}
